#include <stdio.h>
#include <stdlib.h>
#include <apr_pools.h>
#include <apr_hash.h>
#include <apr_thread_mutex.h>
#include <apr_thread_cond.h>
#include "mpi.h"
#include "oom.h"
#include "env.h"
#include "uid.h"
#include "oper.h"

static apr_pool_t *oom_pool;
static apr_hash_t *oom_groups;
static apr_hash_t *oom_comms;
static apr_hash_t *oom_requests;
static apr_hash_t *oom_datatypes;
static apr_hash_t *oom_ops;

static uid_pool context_tag_pool;
static uid_pool group_handle_pool;
static uid_pool comm_handle_pool;
static uid_pool datatype_handle_pool;
static uid_pool request_handle_pool;
static uid_pool op_handle_pool;

static void register_reserved_comms();

void OOM_init()
{
	/** initialize our memory pool and hash tables for storing
	   opaque objects */
	
	apr_pool_create(&oom_pool, NULL);
	apr_pool_tag(oom_pool, "oom");
	oom_groups = apr_hash_make(oom_pool);
	oom_comms = apr_hash_make(oom_pool);
	oom_requests = apr_hash_make(oom_pool);
	oom_datatypes = apr_hash_make(oom_pool);
	oom_ops = apr_hash_make(oom_pool);
	
	/** initialize unique id pools for context tags and opaque object
	   handles */
	
	uid_init(&context_tag_pool, 32);
	uid_init(&group_handle_pool, 32);
	uid_init(&comm_handle_pool, 32);
	uid_init(&datatype_handle_pool, MPI_DATATYPE_BASIC_MAX + 1);
	uid_init(&request_handle_pool, 32);
	uid_init(&op_handle_pool, 32);
	
	/** create and register the reserved communicators */
	
	register_reserved_comms();
	register_reserved_ops();
}

void OOM_finalize()
{
	apr_pool_destroy(oom_pool);
}

MPI_Group OOM_group_create(int *members, int num_members)
{
	MPI_Group *handle;
	Group *group;
	int i;
	
	handle = (MPI_Group *) apr_palloc(oom_pool, sizeof(MPI_Group));
	uid_reserve_next(&group_handle_pool, handle);
	
	group = (Group *) apr_palloc(oom_pool, sizeof(Group));
	group->num_members = num_members;
	group->members = (int *) apr_palloc(oom_pool, sizeof(int) * num_members);
	for (i = 0; i < num_members; i++)
	{
		group->members[i] = members[i];
	}
	
	apr_hash_set(oom_groups, handle, sizeof(MPI_Group), group);
	
	return *handle;
}

void OOM_group_destroy(MPI_Group handle)
{
	apr_hash_set(oom_groups, &handle, sizeof(MPI_Group), NULL);
}

Group *OOM_group_lookup(MPI_Group handle)
{
	return (Group *) apr_hash_get(oom_groups, &handle, sizeof(int));
}

MPI_Comm OOM_comm_create(MPI_Group group, int p2p_context,	int coll_context,
	int source)
{
	MPI_Comm *handle;
	Comm *comm;
	
	handle = (MPI_Comm *) apr_palloc(oom_pool, sizeof(MPI_Comm));
	uid_reserve_next(&comm_handle_pool, handle);
	
	comm = (Comm *) apr_palloc(oom_pool, sizeof(Comm));
	comm->group = group;
	comm->p2p_context = p2p_context;
	comm->coll_context = coll_context;
	comm->source = source;
	comm->is_permanent = 0;
	comm->topology = MPI_UNDEFINED;
	comm->attributes = apr_hash_make(oom_pool);
	
	apr_hash_set(oom_comms, handle, sizeof(MPI_Comm), comm);
	
	return *handle;
}

void OOM_comm_destroy(MPI_Comm handle)
{
	apr_hash_set(oom_comms, &handle, sizeof(MPI_Comm), NULL);
}

Comm *OOM_comm_lookup(MPI_Comm handle)
{
	return (Comm *) apr_hash_get(oom_comms, &handle, sizeof(int));
}

MPI_Datatype OOM_datatype_create()
{
	MPI_Datatype *handle;
	Datatype *datatype;
	
	handle = (MPI_Datatype *) apr_palloc(oom_pool, sizeof(MPI_Datatype));
	uid_reserve_next(&datatype_handle_pool, handle);
	
	datatype = (Datatype *) apr_palloc(oom_pool, sizeof(Datatype));

	datatype->num_elements = 0;
    datatype->typemap = NULL;
    datatype->lb = NULL;
    datatype->ub = NULL;
    datatype->extent = (MPI_Aint) 0;
    datatype->size = 0;
    datatype->committed = 0;
	
	apr_hash_set(oom_datatypes, handle, sizeof(MPI_Datatype), datatype);
	
	return *handle;
}

void OOM_datatype_destroy(MPI_Datatype handle)
{
	apr_hash_set(oom_datatypes, &handle, sizeof(MPI_Datatype), NULL);
}

Datatype *OOM_datatype_lookup(MPI_Datatype handle)
{
	return (Datatype *) apr_hash_get(oom_datatypes, &handle,
		sizeof(MPI_Datatype));
}

Typemap *OOM_typemap_create()
{
	return (Typemap *) apr_palloc(oom_pool, sizeof(Typemap));
}

MPI_Request OOM_request_create()
{
	Request *request;
	MPI_Request *handle;
	apr_status_t status;
	
	
	handle = (MPI_Request *) apr_palloc(oom_pool, sizeof(MPI_Request));
	uid_reserve_next(&request_handle_pool, handle);
	
	request = (Request *) apr_palloc(oom_pool, sizeof(Request));
	status = apr_thread_mutex_create(&(request->mutex),
		APR_THREAD_MUTEX_NESTED, oom_pool);
	if (status != APR_SUCCESS) return MPI_REQUEST_NULL;
	status = apr_thread_cond_create(&(request->cond), oom_pool);
	if (status != APR_SUCCESS) return MPI_REQUEST_NULL;
	
	apr_hash_set(oom_requests, handle, sizeof(MPI_Request), request);
	
	return *handle;
}

void OOM_request_destroy(MPI_Request handle)
{
	apr_hash_set(oom_requests, &handle, sizeof(MPI_Request), NULL);
}

Request *OOM_request_lookup(MPI_Request handle)
{
	return (Request *) apr_hash_get(oom_requests, &handle,
		sizeof(MPI_Request));
}

MPI_Op OOM_op_create_explicit(MPI_User_function *function, int commute, MPI_Op inhandle)
{
	Op *op;
	MPI_Op *handle;
	
	handle = (MPI_Op *) apr_palloc(oom_pool, sizeof(MPI_Op));
	*handle = inhandle;
	
	op = (Op *) apr_palloc(oom_pool, sizeof(Op));
	op->function = function;
	op->commute = commute;
	
	apr_hash_set(oom_ops, handle, sizeof(MPI_Op), op);
	
	return *handle;
}

MPI_Op OOM_op_create(MPI_User_function *function, int commute)
{
	MPI_Op handle;
	
	uid_reserve_next(&op_handle_pool, &handle);
	return OOM_op_create_explicit(function,commute,handle);
}

void OOM_op_destroy(MPI_Op handle)
{
	apr_hash_set(oom_ops, &handle, sizeof(MPI_Op), NULL);
}

Op *OOM_op_lookup(MPI_Op handle)
{
	return (Op *) apr_hash_get(oom_ops, &handle, sizeof(MPI_Op));
}

MPI_Errhandler OOM_errhandler_create()
{
	/* stub */
	return MPI_ERRHANDLER_NULL;
}

void OOM_errhandler_destroy(MPI_Errhandler handle)
{
	/* stub */
}

Errhandler *OOM_errhandler_lookup(MPI_Errhandler handle)
{
	/* stub */
	return NULL;
}

uid_pool *OOM_get_context_tag_pool()
{
	return &context_tag_pool;
}

void *OOM_palloc(apr_size_t size)
{
	return apr_palloc(oom_pool, size);
}

static void register_reserved_comms()
{
	MPI_Group *group_handle;
	Group *group;
	MPI_Comm *comm_handle;
	Comm *comm;
	int i;
	
	/** create the group used by MPI_COMM_WORLD */
	
	group_handle = (int *) apr_palloc(oom_pool, sizeof(MPI_Group));
	uid_reserve_next(&group_handle_pool, group_handle);
	
	group = (Group *) apr_palloc(oom_pool, sizeof(Group));
	group->num_members = get_world_size();
	group->members = (int *) apr_palloc(oom_pool,
		sizeof(int) * group->num_members);
	for (i = 0; i < group->num_members; i++)
	{
		group->members[i] = i;
	}
	group->is_permanent = 1;
	apr_hash_set(oom_groups, group_handle, sizeof(MPI_Group), group);
	
	/* create MPI_COMM_WORLD */
	
	comm_handle = (int *) apr_palloc(oom_pool, sizeof(MPI_Comm));
	*comm_handle = MPI_COMM_WORLD;
	
	comm = (Comm *) apr_palloc(oom_pool, sizeof(Comm));
	comm->group = *group_handle;
	uid_reserve_next(&context_tag_pool, &(comm->p2p_context));
	uid_reserve_next(&context_tag_pool, &(comm->coll_context));
	comm->source = get_world_rank();
	comm->is_permanent = 1;
	apr_hash_set(oom_comms, comm_handle, sizeof(MPI_Comm), comm);
	
	/** create the group used by MPI_COMM_SELF */
	
	group_handle = (int *) apr_palloc(oom_pool, sizeof(MPI_Group));
	uid_reserve_next(&group_handle_pool, group_handle);
	
	group = (Group *) apr_palloc(oom_pool, sizeof(Group));
	group->num_members = 1;
	group->members = (int *) apr_palloc(oom_pool, sizeof(int));
	group->members[0] = get_world_rank();
	group->is_permanent = 1;
	apr_hash_set(oom_groups, group_handle, sizeof(MPI_Group), group);
	
	/** create MPI_COMM_SELF */
	
	comm_handle = (int *) apr_palloc(oom_pool, sizeof(MPI_Comm));
	*comm_handle = MPI_COMM_SELF;
	
	comm = (Comm *) apr_palloc(oom_pool, sizeof(Comm));
	comm->group = *group_handle;
	uid_reserve_next(&context_tag_pool, &(comm->p2p_context));
	uid_reserve_next(&context_tag_pool, &(comm->coll_context));
	comm->source = get_world_rank();
	comm->is_permanent = 1;
	apr_hash_set(oom_comms, comm_handle, sizeof(MPI_Comm), comm);
}
